import pinecone
from langchain_community.vectorstores.pinecone import Pinecone
from langchain_openai import OpenAIEmbeddings

from backend.app.core.conf import settings


class DbConnection(object):
    """向量数据库连接"""
    # _pinecone_api_key = "c14e2b89-f293-43d6-bb48-f8c912df30d7"
    # _pinecone_env = "gcp-starter"
    #
    _index_map = dict()
    _embeddings = None
    _document_map = dict()

    def __init__(self):
        self._pinecone_api_key = settings.DB_PINECONE_API_KEY
        self._pinecone_env = settings.DB_PINECONE_ENV

    def open(self):
        """
        Notice:方便单独运行不需要初始化框架
        :return:
        """
        pinecone.init(
            api_key=self._pinecone_api_key,
            environment=self._pinecone_env
        )
        self._embeddings = OpenAIEmbeddings()

    def close(self):
        return

    def get_rds(self, index_name: str = "avl", namespace: str = "") -> Pinecone:
        if self._index_map and index_name in self._index_map:
            return self._index_map[index_name]

        conn = Pinecone.from_existing_index(
            embedding=self._embeddings,
            index_name=index_name,
            namespace=namespace,
        )
        self._index_map[index_name] = conn
        return conn

    def get_embeddings(self):
        return self._embeddings


# 初始化数据库连接（实现单例模式）
rdb_conn = DbConnection()
